package doubao;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class DouBaoAssistant {
    private static final String API_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions";
    private static final String API_KEY = "60bcbd6a-0c23-4db8-8245-80383e9f9e67";//换成你自己的APIKey
    private static final String MODEL_NAME = "doubao-pro-32k-241215"; //模型名称

    public static void GetAssistantResponse(){
        List<JsonObject> messages = new ArrayList<>();
        JsonObject systemMessage = new JsonObject();
        systemMessage.addProperty("role", "system");
        systemMessage.addProperty("content", "你是豆包，是由字节跳动开发的 AI 人工智能助手");
        messages.add(systemMessage);

        while (true) {
            System.out.print("你（输入“退出”退出豆包）：");
            String userInput = new java.util.Scanner(System.in).nextLine();
            if ("退出".equalsIgnoreCase(userInput)) {
                break;
            }

            JsonObject userMessage = new JsonObject();
            userMessage.addProperty("role", "user");
            userMessage.addProperty("content", userInput);
            messages.add(userMessage);

            String response = sendRequest(messages);

            JsonObject assistantMessage = new JsonObject();
            assistantMessage.addProperty("role", "assistant");
            assistantMessage.addProperty("content", response);
            messages.add(assistantMessage);

            System.out.println("豆包： " + response);
        }
    }

    private static String sendRequest(List<JsonObject> messages) {
        HttpClient client = HttpClients.createDefault();
        HttpPost post = new HttpPost(API_URL);
        post.setHeader("Authorization", "Bearer " + API_KEY);
        post.setHeader("Content-Type", "application/json");

        JsonObject requestBody = new JsonObject();
        requestBody.addProperty("model", MODEL_NAME);
        requestBody.add("messages", new Gson().toJsonTree(messages));

        post.setEntity(new StringEntity(requestBody.toString(), "UTF-8"));

        String response = null;
        try {
            response = EntityUtils.toString(client.execute(post).getEntity());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        JsonObject jsonResponse = new Gson().fromJson(response, JsonObject.class);

        return jsonResponse.getAsJsonArray("choices").get(0).getAsJsonObject().get("message").getAsJsonObject().get("content").getAsString();
    }
}